'''
 * @Author: Benjay·Shaw
 * @Date: 2024-10-30 17:48:20
 * @LastEditors: Benjay·Shaw
 * @LastEditTime: 2024-10-31 14:47:27
 * @Description: 推理公用模块
'''
import paddle
import math
import warnings
from functools import partial
from attention import *
try:
    from osgeo import gdal
    from osgeo import ogr
    from osgeo import osr
except ImportError:
    import gdal
    import ogr
    import osr



class SPPBlock(paddle.nn.Layer):

    def __init__(self, in_channels):
        super(SPPBlock, self).__init__()
        self.pool1 = paddle.nn.MaxPool2D(kernel_size=(2, 2), stride=2)
        self.pool2 = paddle.nn.MaxPool2D(kernel_size=(3, 3), stride=3)
        self.pool3 = paddle.nn.MaxPool2D(kernel_size=(5, 5), stride=5)
        self.conv = paddle.nn.Conv2D(in_channels=in_channels, out_channels=
            1, kernel_size=1, padding=0)

    def forward(self, x):
        in_channels, h, w = x.shape[1], x.shape[2], x.shape[3]
        layer1 = paddle.nn.functional.upsample(x=self.conv(self.pool1(x)),
            size=(h, w), mode='bilinear')
        layer2 = paddle.nn.functional.upsample(x=self.conv(self.pool2(x)),
            size=(h, w), mode='bilinear')
        layer3 = paddle.nn.functional.upsample(x=self.conv(self.pool3(x)),
            size=(h, w), mode='bilinear')
        out = paddle.concat(x=[layer1, layer2, layer3, x], axis=1)
        return out


class DBlockMoreDilate(paddle.nn.Layer):

    def __init__(self, channel):
        super(DBlockMoreDilate, self).__init__()
        self.dilate1 = paddle.nn.Conv2D(in_channels=channel, out_channels=
            channel, kernel_size=3, dilation=1, padding=1)
        self.dilate2 = paddle.nn.Conv2D(in_channels=channel, out_channels=
            channel, kernel_size=3, dilation=2, padding=2)
        self.dilate3 = paddle.nn.Conv2D(in_channels=channel, out_channels=
            channel, kernel_size=3, dilation=4, padding=4)
        self.relu = paddle.nn.ReLU()
        for m in self.sublayers():
            if isinstance(m, paddle.nn.Conv2D) or isinstance(m, paddle.nn.
                Conv2DTranspose):
                if m.bias is not None:
                    m.bias.data.zero_()

    def forward(self, x):
        dilate1_out = self.relu(self.dilate1(x))
        dilate2_out = self.relu(self.dilate2(dilate1_out))
        dilate3_out = self.relu(self.dilate3(dilate2_out))
        out = x + dilate1_out + dilate2_out + dilate3_out
        return out


class DecoderBlock(paddle.nn.Layer):

    def __init__(self, in_channels, n_filters):
        super(DecoderBlock, self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels=in_channels, out_channels
            =in_channels // 4, kernel_size=1)
        self.norm1 = paddle.nn.BatchNorm2D(num_features=in_channels // 4)
        self.relu1 = paddle.nn.ReLU()
        self.de_conv2 = paddle.nn.Conv2DTranspose(in_channels=in_channels //
            4, out_channels=in_channels // 4, kernel_size=3, stride=2,
            padding=1, output_padding=1)
        self.norm2 = paddle.nn.BatchNorm2D(num_features=in_channels // 4)
        self.relu2 = paddle.nn.ReLU()
        self.conv3 = paddle.nn.Conv2D(in_channels=in_channels // 4,
            out_channels=n_filters, kernel_size=1)
        self.norm3 = paddle.nn.BatchNorm2D(num_features=n_filters)
        self.relu3 = paddle.nn.ReLU()

    def forward(self, x):
        x = self.conv1(x)
        x = self.norm1(x)
        x = self.relu1(x)
        x = self.de_conv2(x)
        x = self.norm2(x)
        x = self.relu2(x)
        x = self.conv3(x)
        x = self.norm3(x)
        x = self.relu3(x)
        return x